import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from enum import Enum

# 飞行模式枚举
class FlightMode(Enum):
    TAKEOFF = 0
    WAYPOINT = 1
    RETURN_HOME = 2
    LANDING = 3

# 测试模式枚举
class TestMode(Enum):
    TAKEOFF_TEST = 0
    LANDING_TEST = 1
    WAYPOINT_TEST = 2
    RETURN_HOME_TEST = 3
    COMPLETE_TEST = 4

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长
SIM_TIME = 30      # 仿真时长（秒）
INERTIA = np.array([0.1, 0.1, 0.1])  # 转动惯量

# 飞行参数
TAKEOFF_HEIGHT = 12.0  # 起飞高度
LANDING_SPEED = 0.5   # 降落速度
WAYPOINT_RADIUS = 0.1 # 到达航路点的判定半径
HOME_POSITION = np.array([0.0, 0.0, 0.0])  # 返航点位置

# 航路点列表
WAYPOINTS = [
    np.array([2.0, 2.0, 2.0]),
    np.array([2.0, -2.0, 2.0]),
    np.array([-2.0, -2.0, 2.0]),
    np.array([-2.0, 2.0, 2.0])
]

# 四旋翼几何参数
ARM_LENGTH = 0.2   # 旋翼臂长
ROTOR_RADIUS = 0.05  # 旋翼半径
AXIS_LENGTH = 0.3  # 姿态指示器长度

# 测试参数
TEST_COMPLETION_TIME = 5.0  # 每个测试模式完成后等待时间
test_mode = TestMode.WAYPOINT_TEST  # 当前测试模式
test_completion_time = 0  # 测试完成计时器

# 速度控制参数
MAX_SPEED = 2.0    # 最大飞行速度 (m/s)
APPROACH_DISTANCE = 1.0  # 接近航路点时开始减速的距离

def create_quadcopter_vertices(pos, att):
    """创建四旋翼的顶点"""
    # 创建基础顶点（未旋转）
    vertices = np.array([
        [ARM_LENGTH, 0, 0],    # 右前
        [0, ARM_LENGTH, 0],    # 前
        [-ARM_LENGTH, 0, 0],   # 左前
        [0, -ARM_LENGTH, 0],   # 后
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    vertices = vertices @ R.T
    
    # 添加位置偏移
    vertices += pos
    
    return vertices

def create_attitude_indicators(pos, att):
    """创建姿态指示器"""
    # 创建基础坐标轴
    axes = np.array([
        [AXIS_LENGTH, 0, 0],    # X轴
        [0, AXIS_LENGTH, 0],    # Y轴
        [0, 0, AXIS_LENGTH],    # Z轴
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    axes = axes @ R.T
    
    # 添加位置偏移
    axes += pos
    
    return axes

def create_quadcopter_faces(vertices):
    """创建四旋翼的面"""
    # 创建中心点
    center = np.mean(vertices, axis=0)
    
    # 创建连接中心点和各个顶点的三角形
    faces = []
    for i in range(4):
        faces.append([center, vertices[i], vertices[(i+1)%4]])
    
    return faces

# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 状态变量
pos = np.array([0.0, 0.0, 0.0])   # 位置 x, y, z
vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
ang_vel = np.array([0.0, 0.0, 0.0])  # 角速度

# 飞行模式
current_mode = FlightMode.TAKEOFF
current_waypoint_index = 0  # 提前声明

# 位置PID控制器
pid_x = PID(2.0, 0.0, 1.0)
pid_y = PID(2.0, 0.0, 1.0)
pid_z = PID(10.0, 2.0, 5.0)

# 速度PID控制器
pid_vx = PID(2.0, 0.0, 0.5)
pid_vy = PID(2.0, 0.0, 0.5)
pid_vz = PID(2.0, 0.0, 0.5)

# 姿态控制器
pid_roll = PID(2.0, 0.0, 0.5)
pid_pitch = PID(2.0, 0.0, 0.5)
pid_yaw = PID(2.0, 0.0, 0.5)

# 确保起飞测试时有目标高度
if test_mode == TestMode.TAKEOFF_TEST:
    pid_z.set_setpoint(TAKEOFF_HEIGHT)
# 确保航路点测试时有目标航点
if test_mode == TestMode.WAYPOINT_TEST:
    current_mode = FlightMode.WAYPOINT
    current_waypoint_index = 0
    pid_x.set_setpoint(WAYPOINTS[0][0])
    pid_y.set_setpoint(WAYPOINTS[0][1])
    pid_z.set_setpoint(WAYPOINTS[0][2])

# 存储历史数据
positions = []
attitudes = []
flight_modes = []
velocities = []

# 打开文件以记录位置和姿态数据
with open("quadcopter_data.txt", "w") as file:
    file.write("Time (s),Longitude,Latitude,Altitude,Roll (deg),Pitch (deg),Yaw (deg)\n")

    def calculate_velocity_setpoint(target_pos, current_pos):
        """计算速度设定点，考虑最大速度和接近目标时的减速"""
        # 计算到目标的向量
        direction = target_pos - current_pos
        distance = np.linalg.norm(direction)
        
        if distance < 0.01:  # 防止除以零
            return np.zeros(3)
        
        # 单位方向向量
        direction_unit = direction / distance
        
        # 根据距离计算目标速度
        if distance > APPROACH_DISTANCE:
            # 远距离时使用最大速度
            return direction_unit * MAX_SPEED
        else:
            # 接近目标时线性减速
            speed_factor = distance / APPROACH_DISTANCE
            return direction_unit * MAX_SPEED * speed_factor

    def update_test_mode():
        global test_mode, current_mode, test_completion_time, pos, vel, att, ang_vel, current_waypoint_index
        
        if test_completion_time > 0:
            test_completion_time -= DT
            return
        
        if test_mode == TestMode.TAKEOFF_TEST:
            if current_mode == FlightMode.TAKEOFF and pos[2] >= TAKEOFF_HEIGHT:
                test_mode = TestMode.LANDING_TEST
                current_mode = FlightMode.LANDING
                pid_z.set_setpoint(0)
                test_completion_time = TEST_COMPLETION_TIME
        
        elif test_mode == TestMode.LANDING_TEST:
            if current_mode == FlightMode.LANDING and pos[2] <= 0.1:
                test_mode = TestMode.WAYPOINT_TEST
                current_mode = FlightMode.TAKEOFF
                current_waypoint_index = 0
                pid_x.set_setpoint(WAYPOINTS[0][0])
                pid_y.set_setpoint(WAYPOINTS[0][1])
                pid_z.set_setpoint(WAYPOINTS[0][2])
                test_completion_time = TEST_COMPLETION_TIME
        
        elif test_mode == TestMode.WAYPOINT_TEST:
            if current_mode == FlightMode.WAYPOINT and current_waypoint_index >= len(WAYPOINTS):
                test_mode = TestMode.RETURN_HOME_TEST
                current_mode = FlightMode.RETURN_HOME
                pid_x.set_setpoint(HOME_POSITION[0])
                pid_y.set_setpoint(HOME_POSITION[1])
                pid_z.set_setpoint(HOME_POSITION[2])
                test_completion_time = TEST_COMPLETION_TIME
        
        elif test_mode == TestMode.RETURN_HOME_TEST:
            if current_mode == FlightMode.RETURN_HOME and np.linalg.norm(pos - HOME_POSITION) < WAYPOINT_RADIUS:
                test_mode = TestMode.COMPLETE_TEST
                current_mode = FlightMode.LANDING
                pid_z.set_setpoint(0)
                test_completion_time = TEST_COMPLETION_TIME

    def update_flight_mode():
        global current_mode, current_waypoint_index
        
        if current_mode == FlightMode.TAKEOFF:
            if pos[2] >= TAKEOFF_HEIGHT:
                if test_mode == TestMode.TAKEOFF_TEST:
                    return
                current_mode = FlightMode.WAYPOINT
                current_waypoint_index = 0
                pid_x.set_setpoint(WAYPOINTS[0][0])
                pid_y.set_setpoint(WAYPOINTS[0][1])
                pid_z.set_setpoint(WAYPOINTS[0][2])
        
        elif current_mode == FlightMode.WAYPOINT:
            current_waypoint = WAYPOINTS[current_waypoint_index]
            distance = np.linalg.norm(pos - current_waypoint)
            
            if distance < WAYPOINT_RADIUS:
                current_waypoint_index += 1
                if current_waypoint_index >= len(WAYPOINTS):
                    if test_mode == TestMode.WAYPOINT_TEST:
                        return
                    current_mode = FlightMode.RETURN_HOME
                    pid_x.set_setpoint(HOME_POSITION[0])
                    pid_y.set_setpoint(HOME_POSITION[1])
                    pid_z.set_setpoint(HOME_POSITION[2])
                else:
                    next_waypoint = WAYPOINTS[current_waypoint_index]
                    pid_x.set_setpoint(next_waypoint[0])
                    pid_y.set_setpoint(next_waypoint[1])
                    pid_z.set_setpoint(next_waypoint[2])
        
        elif current_mode == FlightMode.RETURN_HOME:
            distance = np.linalg.norm(pos - HOME_POSITION)
            if distance < WAYPOINT_RADIUS:
                if test_mode == TestMode.RETURN_HOME_TEST:
                    return
                current_mode = FlightMode.LANDING
                pid_z.set_setpoint(0)
        
        elif current_mode == FlightMode.LANDING:
            # 让 setpoint 以最大下降速度递减
            new_setpoint = max(pid_z.setpoint - LANDING_SPEED * DT, 0)
            pid_z.set_setpoint(new_setpoint)
            if pos[2] <= 0.1:
                if test_mode == TestMode.COMPLETE_TEST:
                    return
                current_mode = FlightMode.TAKEOFF

    # 创建图形
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')

    # 设置坐标轴范围
    ax.set_xlim([-3, 3])
    ax.set_ylim([-3, 3])
    ax.set_zlim([-3, 3])
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # 绘制航路点
    for waypoint in WAYPOINTS:
        ax.scatter(waypoint[0], waypoint[1], waypoint[2], color='green', marker='^')

    # 绘制返航点
    ax.scatter(HOME_POSITION[0], HOME_POSITION[1], HOME_POSITION[2], color='red', label='Home')

    # 创建四旋翼多边形
    vertices = create_quadcopter_vertices(pos, att)
    faces = create_quadcopter_faces(vertices)
    quadcopter = Poly3DCollection(faces, alpha=0.5, facecolor='blue', edgecolor='black')
    ax.add_collection3d(quadcopter)

    # 创建姿态指示器
    axes = create_attitude_indicators(pos, att)
    x_axis, = ax.plot([pos[0], axes[0,0]], [pos[1], axes[0,1]], [pos[2], axes[0,2]], 'r-', linewidth=2)
    y_axis, = ax.plot([pos[0], axes[1,0]], [pos[1], axes[1,1]], [pos[2], axes[1,2]], 'g-', linewidth=2)
    z_axis, = ax.plot([pos[0], axes[2,0]], [pos[1], axes[2,1]], [pos[2], axes[2,2]], 'b-', linewidth=2)

    # 创建轨迹线
    trajectory_line, = ax.plot([], [], [], 'b-', alpha=0.5)

    # 创建速度向量
    velocity_vector, = ax.plot([pos[0], pos[0]+vel[0]*0.2], [pos[1], pos[1]+vel[1]*0.2], [pos[2], pos[2]+vel[2]*0.2], 'y-', linewidth=2)

    # 添加模式文本
    mode_text = ax.text2D(0.02, 0.95, '', transform=ax.transAxes)
    speed_text = ax.text2D(0.02, 0.90, '', transform=ax.transAxes)

    def update(frame):
        global pos, vel, att, ang_vel
        
        # 更新测试模式
        update_test_mode()
        
        # 更新飞行模式
        update_flight_mode()
        
        # 计算速度设定点
        if current_mode == FlightMode.TAKEOFF:
            # 起飞时主要控制高度速度
            vx_setpoint = 0
            vy_setpoint = 0
            vz_setpoint = min(MAX_SPEED, (TAKEOFF_HEIGHT - pos[2]) / 2)  # 逐渐增加上升速度
        
        elif current_mode == FlightMode.WAYPOINT:
            # 航路点模式下计算速度设定点
            target_pos = WAYPOINTS[current_waypoint_index]
            velocity_setpoint = calculate_velocity_setpoint(target_pos, pos)
            vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
        
        elif current_mode == FlightMode.RETURN_HOME:
            # 返航模式下计算速度设定点
            velocity_setpoint = calculate_velocity_setpoint(HOME_POSITION, pos)
            vx_setpoint, vy_setpoint, vz_setpoint = velocity_setpoint
        
        elif current_mode == FlightMode.LANDING:
            # 降落时主要控制下降速度
            vx_setpoint = 0
            vy_setpoint = 0
            vz_setpoint = -min(LANDING_SPEED, (pos[2] - pid_z.setpoint) / 2)  # 逐渐增加下降速度
        
        # 设置速度PID控制器的设定点
        pid_vx.set_setpoint(vx_setpoint)
        pid_vy.set_setpoint(vy_setpoint)
        pid_vz.set_setpoint(vz_setpoint)
        
        # 速度PID控制器输出
        fx = pid_vx.update(vel[0], DT)
        fy = pid_vy.update(vel[1], DT)
        fz = pid_vz.update(vel[2], DT) + MASS * GRAVITY

        # 计算姿态控制力矩
        tau_roll = pid_roll.update(att[0], DT)
        tau_pitch = pid_pitch.update(att[1], DT)
        tau_yaw = pid_yaw.update(att[2], DT)

        # 合力
        force = np.array([fx, fy, fz])
        torque = np.array([tau_roll, tau_pitch, tau_yaw])

        # 加速度
        acc = force / MASS - np.array([0, 0, GRAVITY])
        ang_acc = torque / INERTIA

        # 更新速度和位置
        vel += acc * DT
        pos += vel * DT
        
        # 防止降到地面以下
        if pos[2] < 0:
            pos[2] = 0
            vel[2] = 0

        # 更新角速度和姿态
        ang_vel += ang_acc * DT
        att += ang_vel * DT

        positions.append(pos.copy())
        attitudes.append(att.copy())
        flight_modes.append(current_mode)
        velocities.append(vel.copy())

        # 写入数据到文件
        time_elapsed = frame * DT
        longitude = -116.6847222 + pos[0] * 0.0001  # 假定的经度转换
        latitude = 32.72722222 + pos[1] * 0.0001  # 假定的纬度转换
        altitude = pos[2]
        roll_deg = np.degrees(att[0])  # 将弧度转换为度
        pitch_deg = np.degrees(att[1])  # 将弧度转换为度
        yaw_deg = np.degrees(att[2])  # 将弧度转换为度
        
        file.write(f"{time_elapsed:.2f},{longitude:.8f},{latitude:.8f},{altitude:.2f},{roll_deg:.5f},{pitch_deg:.5f},{yaw_deg:.5f}\n")

        # 更新四旋翼位置和姿态
        vertices = create_quadcopter_vertices(pos, att)
        faces = create_quadcopter_faces(vertices)
        quadcopter.set_verts(faces)

        # 更新姿态指示器
        axes = create_attitude_indicators(pos, att)
        x_axis.set_data([pos[0], axes[0,0]], [pos[1], axes[0,1]])
        x_axis.set_3d_properties([pos[2], axes[0,2]])
        y_axis.set_data([pos[0], axes[1,0]], [pos[1], axes[1,1]])
        y_axis.set_3d_properties([pos[2], axes[1,2]])
        z_axis.set_data([pos[0], axes[2,0]], [pos[1], axes[2,1]])
        z_axis.set_3d_properties([pos[2], axes[2,2]])

        # 更新轨迹线
        positions_array = np.array(positions)
        trajectory_line.set_data(positions_array[:, 0], positions_array[:, 1])
        trajectory_line.set_3d_properties(positions_array[:, 2])
        
        # 更新速度向量
        velocity_magnitude = np.linalg.norm(vel)
        velocity_vector.set_data([pos[0], pos[0]+vel[0]*0.2], [pos[1], pos[1]+vel[1]*0.2])
        velocity_vector.set_3d_properties([pos[2], pos[2]+vel[2]*0.2])

        # 更新模式文本
        mode_text.set_text(f'Test Mode: {test_mode.name}\nFlight Mode: {current_mode.name}')
        speed_text.set_text(f'Speed: {velocity_magnitude:.2f} m/s')

        return quadcopter, x_axis, y_axis, z_axis, trajectory_line, velocity_vector, mode_text, speed_text

    # 创建动画
    ani = FuncAnimation(fig, update, frames=int(SIM_TIME/DT), interval=DT*1000, blit=True)
    plt.legend()
    plt.tight_layout()
    plt.show()
